import os
from PIL import Image
from torchvision import transforms
import torch
from tqdm import tqdm

# 配置路径
hr_dir = 'SR_Datasets/DIV2K/DIV2K_train_HR'
lr_dir = 'SR_Datasets/DIV2K/DIV2K_train_LR_bicubic/X4'
output_hr_dir = 'SR_Datasets/pt/HR'
output_lr_dir = 'SR_Datasets/pt/LR'

# 创建输出目录
os.makedirs(output_hr_dir, exist_ok=True)
os.makedirs(output_lr_dir, exist_ok=True)

# 转换设置
hr_size = (512, 512)
lr_size = (128, 128)
to_tensor = transforms.ToTensor()

# 获取文件名
image_files = sorted([
    f for f in os.listdir(hr_dir)
    if f.endswith('.jpg') or f.endswith('.png')
])

# 处理并保存
for name in tqdm(image_files, desc="Converting to .pt"):
    hr_path = os.path.join(hr_dir, name)
    lr_path = os.path.join(lr_dir, name.replace('.png', 'x4.png').replace('.jpg', 'x4.png'))

    # 加载图像
    hr_img = Image.open(hr_path).convert('RGB').resize(hr_size, Image.LANCZOS)
    lr_img = Image.open(lr_path).convert('RGB').resize(lr_size, Image.LANCZOS)

    # 转 tensor
    hr_tensor = to_tensor(hr_img)
    lr_tensor = to_tensor(lr_img)

    # 保存 .pt 文件
    torch.save(hr_tensor, os.path.join(output_hr_dir, name + '.pt'))
    torch.save(lr_tensor, os.path.join(output_lr_dir, name + '.pt'))
